#! /usr/bin/python
# -*- coding: utf-8 -*-
## Author: Dean Chen, deanchen@juniper.net

from ftplib import FTP
import os
import sys
import glob
import getopt
import threading
import telnetlib
import time
import re
import pexpect
from pexpect import *
import pxssh

class build_install:

		#######################################################
		#You can modify below variate according to your demand
		#######################################################

        device_name               = ''
        build_name                = ''
        build_location            = ''
        #build_save_location       = '/var/tmp'
        build_save_location       = '/var/crash/corefiles/'
        #default_location_vsrx     = '/homes/ssd-builder/ngsrx-yocto-daily/RELEASE_174_THROTTLE//LATEST/srx-mr/ship/cli'
        default_location_vsrx     = '/homes/ssd-builder/ngsrx-yocto-daily/DEV_COMMON_BRANCH/LATEST/srx-mr/ship/cli'
        #default_location_vsrx     = '/homes/ssd-builder/ngsrx-yocto-daily/JUNOS_174_R1_BRANCH/LATEST/srx-mr/ship/cli'
        default_location_summit     = '/homes/ssd-builder/ngsrx-yocto-daily/JUNOS_174_R1_BRANCH/LATEST/srx-summit/ship/cli'
        #default_location_summit     = '/homes/ssd-builder/ngsrx-yocto-daily/RELEASE_174_THROTTLE/LATEST/srx-summit/ship/cli'
        #default_location_summit     = '/homes/ssd-builder/ngsrx-yocto-daily/DEV_COMMON_BRANCH/LATEST/srx-summit/ship/cli'
        default_location_GW     = '/homes/ssd-builder/ngsrx-yocto-daily/DEV_COMMON_BRANCH/LATEST/srx-mr/ship/cli'
        default_location_Forge    = '/homes/ssd-builder/ngsrx-yocto-daily/DEV_COMMON_BRANCH/LATEST/srx-700e/ship/cli'
        default_location_HE_Seige = '/volume/build/junos/dev_common/17.4/development/current/ship'
#        default_location_HE_Seige = '/homes/slt-builder/sandboxes/DEV_X_151_X49_BRANCH/current/ship/'
        build_type                = ''
		
		#######################################################
		#You can modify below variate according to your demand
		#######################################################




        def get_device_model(self,device_name):
				global model
				tn = telnetlib.Telnet(device_name)
				tn.set_debuglevel(0)
				tn.read_until("login: ")
				tn.write("regress" + "\n")
				tn.read_until("Password:")
				tn.write("MaRtInI" + "\n")
				tn.write("cli" + "\n")
				tn.write("show version|match Model" + "\n")
				tn.write("exit\n")
				time.sleep(3)
				tn.write("exit\n")
				model = re.findall(r"Model: (.*)",tn.read_all())
				print '######################################################'
				print '###### '+device_name+' model is : '+ model[0]+' ######'
				print '######################################################'
		
        def download_build(self,device_name,build_name,build_location,build_save_location):


				self.get_build_info(device_name,build_location,build_name)

				ftp=FTP()
				build=self.build_location+ "/" +self.build_name
				ftp.set_debuglevel(1)
				ftp.connect(device_name,'21')
				ftp.login('root','Embe1mpls')
				print ftp.getwelcome()
				ftp.cwd(build_save_location)
				bufsize = 1024
				file_handler = open(build,'rb')
				lsize = os.path.getsize(build)
				print '###################################################'
				print '###### The build\'s size is :'+ str(lsize)+' ######'
				print '###################################################'
				ftp.storbinary('STOR %s' % os.path.basename(build),file_handler,bufsize)
				ftp.set_debuglevel(0)
				ftp.quit
				print '#########################################################'
				print "###### Finished build download for "+device_name+' ######'
				print '#########################################################'
				
        def install_build(self,device_name,build_name,build_location,build_save_location):

				self.download_build(device_name,build_name,build_location,build_save_location)

				install_build = "request system software add " +build_save_location+ "/" + self.build_name +" no-copy no-validate reboot"
				tn = telnetlib.Telnet(device_name)
				tn.set_debuglevel(1)
				tn.read_until("login: ")
				tn.write("regress" + "\n")
				tn.read_until("Password:")
				tn.write("MaRtInI" + "\n")
				tn.write("cli" + "\n")
				tn.write(install_build+"\n")
				tn.write("exit\n")
				print tn.read_all()



		
        def get_build_info(self,device_name,build_location,build_name):

				self.get_device_model(device_name)
				global build_list 

				if build_location == '':
						if model[0].strip() == 'vsrx':
								self.build_location = self.default_location_vsrx
								build_list = glob.glob(self.default_location_vsrx +"/"+'*.tgz')
						
						elif model[0].strip() == 'srx4200':
								self.build_location = self.default_location_GW
								build_list = glob.glob(self.default_location_GW +"/"+'*.tgz')
						
						elif model[0].strip() == 'srx4600':
								self.build_location = self.default_location_summit
								build_list = glob.glob(self.default_location_summit +"/"+'*.tgz')
						
						elif model[0].strip() == 'srx1500':
								self.build_location = self.default_location_Forge
								build_list = glob.glob(self.default_location_Forge +"/"+'junos-srxentedge*.tgz')
						
						elif (model[0].strip() == 'srx550m') or (model[0].strip() =='srx345') or (model[0].strip() =='srx320'):
								self.build_location = self.default_location_HE_Seige
								build_list = glob.glob(self.default_location_HE_Seige +"/"+'junos-srxsme*.tgz')
						
						elif (model[0].strip() == 'srx5600') or (model[0].strip() =='srx5800'):
								self.build_location = self.default_location_HE_Seige
								build_list = glob.glob(self.default_location_HE_Seige +"/"+'junos-srx5000*.tgz')

						for build in build_list:
								#select the iamge which we need
								if self.build_type == 'occam':
										if ('x86' in build):
												self.build_name = build.split('/')[-1]
												print '#####################################################################################'
												print '###### Build name for '+ model[0].strip() +': '+self.build_name+ ' ######'
												print '#####################################################################################'

								elif self.build_type != 'occam':
										if ('fips' not in build) and ('export' not in build) and ('x86' not in build):
										    self.build_name = build.split('/')[-1]
										    print '#####################################################################################'
										    print '###### Build name for '+ model[0].strip() +': '+self.build_name+ ' ######'
										    print '#####################################################################################'
						return self.build_name



				else:
						if build_name == '':
								if (model[0].strip() == 'vsrx') or (model[0].strip() == 'srx4200') or (model[0].strip() == 'srx4600') or (model[0].strip() == 'srx1500') :
										build_list = glob.glob(self.build_location +"/"+'*.tgz')

								elif model[0].strip() == 'srx550m' or 'srx345' or 'srx320':
										build_list = glob.glob(self.build_ocation +"/"+'junos-srxsme*.tgz')

								elif model[0].strip() == 'srx5600' or 'srx5800':
										build_list = glob.glob(self.build_location +"/"+'junos-srx5000*.tgz')

								for build in build_list:
										#select the image which we need
										if self.build_type == 'occam':
												if ('x86' in build):
														self.build_name = build.split('/')[-1]
														print '#####################################################################################'
														print '###### Build name for '+ model[0].strip() +': '+self.build_name+ ' ######'
														print '#####################################################################################'

										elif self.build_type != 'occam':
												if ('fips' not in build) and ('export' not in build) and ('x86' not in build):
														self.build_name = build.split('/')[-1]
														print '#####################################################################################'
														print '###### Build name for '+ model[0].strip() +': '+self.build_name+ ' ######'
														print '#####################################################################################'

								return self.build_name

						else:
								print '###############################################'
								print '###### Build name : '+self.build_name+ ' ######'
								print '###############################################'
								return self.build_name


class Main_fun:


		def main(self,argv):

				try:
						opts,args = getopt.getopt(sys.argv[1:],"hd:n:l:s:t:",["help"])
				except getopt.GetoptError:
						print "please use './build_install.py -h ' for detail :)"
						sys.exit(2)

				for option,value in opts:
						if option in ["-h","--help"]:
								print '''

                ###### Script support SRX5800/5600/4600/4200/4100/1500/550/345/320/VSRX ######
                ###### For HE and Seige, Please su to regress first, then run the script######


                Usage: ./load_build.py  -d device_name [-n build_name] [-l build_location] [-s build_save_location] [install]

                    -d    devices name which need install build
                    -n    name of the build
                    -l    location of the build
                    -s    location where build save on devices
                    -t    specify 'occam' to load OCCAM iamge

                Examples:

                    ./load_build.py  -d 'slt101-tiger slt101-lion' -n junos-srxmr-15.1-2016-09-08.0_DEV_X_151_X49-domestic.tgz -l /homes/ssd-builder/ngsrx-yocto-daily/DEV_X_151_X49_BRANCH/LATEST/srx-mr/ship/cli/ -s /var/crash/corefiles/
                    ./load_build.py  -d 'slt101-tiger slt101-lion'  #just download but not install, and use the default value in the script, that part is user-defined
                    ./load_build.py  -d 'slt101-tiger slt101-lion' install  #download and install
                    ./load_build.py  -d 'slt15-ginseng guo calla shaoer zinnia pdt-berry' #download build for HE/Summit/Forge/GW/Seige in parallel
                    ./load_build.py  -d  slt102-python -t occam     #download OCCAM build 
								'''
								sys.exit(1)

						elif option == '-d':
								build_install.device_name = value

						elif option == '-n':
								build_install.build_name = value

						elif option == '-l':
								build_install.build_location = value

						elif option == '-s':
								build_install.build_save_location = value

						elif option == '-t':
								build_install.build_type = value


				print '################################################'
				print '###### Device List: '+build_install.device_name+' ######'
				print '################################################'
				
				print '###############################################################'
				print '###### Build Save Location: '+build_install.build_save_location+' ######'
				print '###############################################################'

				device_list = build_install.device_name.split()
				thread_list = []

				if len(args):
						for argument in args:
								if argument == 'install':
										print '############################################'
										print '###### Download and Install build now ######'
										print '############################################'
										for device in device_list:
												install_build_thread = threading.Thread(target = build_install().install_build, args = (device,build_install.build_name,build_install.build_location,build_install.build_save_location))
												install_build_thread.start()


				else:
						print '################################'
						print '###### Download Build Now ######'
						print '################################'
						for device in device_list:
								download_build_thread = threading.Thread(target = build_install().download_build,args = (device,build_install.build_name,build_install.build_location,build_install.build_save_location))
								download_build_thread.start()


		def run_with_regress(self):
#				s = pxssh.pxssh()
#				if not s.login('localhost','regress','MaRtInI'):
#						print "switch to regress failed"
#						print str(s)
#				else:
#						print "lalala"
#						print s.sendline('whoami')
#						s.prompt()
#						print s.before
				self.main(sys.argv[1:])
#a = build_install()
#a.run_with_regress()
a = Main_fun()
a.run_with_regress()









